/*
 * Atomic Predicates for Transformers
 * 
 * Copyright (c) 2015 UNIVERSITY OF TEXAS AUSTIN. All rights reserved. Developed
 * by: HONGKUN YANG and SIMON S. LAM http://www.cs.utexas.edu/users/lam/NRL/
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * with the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimers.
 * 
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimers in the documentation
 * and/or other materials provided with the distribution.
 * 
 * 3. Neither the name of the UNIVERSITY OF TEXAS AUSTIN nor the names of the
 * developers may be used to endorse or promote products derived from this
 * Software without specific prior written permission.
 * 
 * 4. Any report or paper describing results derived from using any part of this
 * Software must cite the following publication of the developers: Hongkun Yang
 * and Simon S. Lam, Scalable Verification of Networks With Packet Transformers
 * Using Atomic Predicates, IEEE/ACM Transactions on Networking, October 2017,
 * Volume 25, No. 5, pages 2900-2915 (first published as IEEE Early Access
 * Article, July 2017, Digital Object Identifier: 10.1109/TNET.2017.2720172).
 * 
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH
 * THE SOFTWARE.
 */
package transformer;

import java.io.*;
import java.util.*;

import common.APComputer;
import common.BDDACLWrapper;
import common.ForwardingRule;
import common.FWDAPSet;
import jdd.bdd.BDD;

public class MPLSTable {

    static BDDACLWrapper bddengine;

    static int reserveLabels_bdd;

    static final String recv_port = "recv";

    class table_entry{
        int out_label;
        HashSet<String> ports;

        table_entry(int ol, HashSet<String> ps)
        {
            out_label = ol;
            ports = ps;
        }
    }

    class table_entry_bdd{
        int out_label_bdd;
        HashSet<String> ports;

        table_entry_bdd(int ol_bdd, HashSet<String> ps)
        {
            out_label_bdd = ol_bdd;
            ports = ps;
        }
    }



    public static void set_bdd_engine(BDDACLWrapper baw)
    {
        bddengine = baw;
        reserveLabels_bdd = baw.getMPLSReserveLabels();
    }

    public int num_rules()
    {
        return swap_table.size();
    }

    public int num_bdd_nodes()
    {
        int num = 0;


        for(int key : swap_table_bdd.keySet())
        {
            num += bddengine.getNodeSize(key);
            num += bddengine.getNodeSize(swap_table_bdd.get(key).out_label_bdd);
        }

        return num;
    }

    static public int MPLS_REM = -1;
    static public int MPLS_DSCD = -2;

    HashMap<Integer, table_entry> swap_table;
    HashMap<Integer, table_entry_bdd> swap_table_bdd;
    // for now, FEC is a prefix
    HashMap<ForwardingRule, table_entry> entrance_table;
    HashMap<Integer, table_entry_bdd> entrance_table_bdd;
    int all_encapsulated;

    // for atomic predicates
    // atomic predicate -> transformed predicate
    HashMap<Integer, HashMap<String, Integer>> predicatesMap;
    HashSet<Integer> computedPredicates;
    // transformationMap stores entries for both label switching
    // and encapsulation
    HashMap<Integer, HashMap<String, FWDAPSet>> transformationMap;

    public void setaps(APComputer apc)
    {
        HashSet<Integer> aps = apc.getAllAP();
        transformationMap = new HashMap<Integer, HashMap<String, FWDAPSet>>();

        if(predicatesMap != null)
        {
            for(int predicate : predicatesMap.keySet())
            {
                if(aps.contains(predicate))
                {
                    HashMap<String, Integer> afters = predicatesMap.get(predicate);
                    HashMap<String, FWDAPSet> after_aps = new HashMap<String, FWDAPSet> ();

                    for(String port : afters.keySet())
                    {
                        int ap = afters.get(port);

                        HashSet<Integer> rawset = apc.getAPExpComputed(ap);
                        if(rawset == null)
                        {
                            System.err.println("bdd expression not found!");
                            System.exit(1);
                        }else
                        {
                            FWDAPSet apset = new FWDAPSet(rawset);
                            after_aps.put(port, apset);
                        }

                    }

                    transformationMap.put(predicate, after_aps);

                }
            }
        }

        // take a look
        //for(int ap : transformationMap.keySet())
        //{
        //    System.out.println(ap + "->" + transformationMap.get(ap));
        //}
    }

    public HashSet<Integer> get_all_changed_predicates()
    {
        HashSet<Integer> predicates = new HashSet<Integer>();
        for(HashMap<String, Integer> one_v : predicatesMap.values())
        {
            predicates.addAll(one_v.values());
        }

        return predicates;
    }

    public void updatePredicatesMap(APComputer apc)
    {
        if(predicatesMap == null)
        {
            predicatesMap = new HashMap<Integer, HashMap<String,Integer>> ();
            computedPredicates = new HashSet<Integer> ();
        }

        for(int ap : apc.getAllAP())
        {
            if(!computedPredicates.contains(ap))
            {
                computedPredicates.add(ap);
                if(bddengine.isMPLSPkt(ap))
                {
                    // test ap for label switching
                    HashMap<String, Integer> fwded = ForwardAction(ap);
                    if(!fwded.isEmpty())
                    {
                        predicatesMap.put(ap, fwded);
                        //System.out.println("enter switching");

                    }
                }else
                {
                    // test ap for encapsulation
                    HashMap<String, Integer> encaped = EncapsulateAction(ap);
                    if(!encaped.isEmpty())
                    {
                        predicatesMap.put(ap, encaped);
                        //System.out.println("enter encapsulation");

                    }
                }
            }
        }
    }

    public MPLSTable()
    {
        swap_table = new HashMap<Integer, table_entry>();
        swap_table_bdd = new HashMap<Integer, table_entry_bdd> ();

        entrance_table = new HashMap<ForwardingRule, table_entry> ();
        entrance_table_bdd = new HashMap<Integer, table_entry_bdd> ();
    }

    public HashMap<Integer, table_entry_bdd> copy_swap_table_bdd()
    {
        return new HashMap<Integer, table_entry_bdd> (swap_table_bdd);
    }

    public void add_swap_entry(int in_label, int out_label, HashSet<String> ports)
    {
        swap_table.put(in_label, new table_entry(out_label, ports));
        //System.out.println("add " + in_label + "," + out_label + "," + ports);
    }
    
    public void add_swap_entry_bdd(int in_label, int out_label, HashSet<String> ports)
    {
        table_entry entry = new table_entry(out_label, ports);
        swap_table.put(in_label, new table_entry(out_label, ports));
        
        int in_label_bdd = bddengine.encodeMPLSLabel(in_label);
        table_entry_bdd entry_bdd = convert_to_bdd_entry(entry);
        swap_table_bdd.put(in_label_bdd, entry_bdd);  
        
        // reset
        computedPredicates = null;
        predicatesMap = null;
    }

    public void add_entrance_entry(ForwardingRule fec, int out_label, HashSet<String> ports)
    {
        entrance_table.put(fec, new table_entry(out_label, ports));
    }
    
    public void add_entrance_entry_bdd(ForwardingRule fec, int out_label, HashSet<String> ports)
    {
        table_entry entry = new table_entry(out_label, ports);
        entrance_table.put(fec, entry);
        
        int fec_bdd = bddengine.encodeDstIPPrefix(fec.getdestip(), fec.getprefixlen());
        table_entry_bdd entry_bdd = convert_to_bdd_entry(entry);
        entrance_table_bdd.put(fec_bdd, entry_bdd);
        
        all_encapsulated = bddengine.getBDD().orTo(all_encapsulated, fec_bdd);
        
        //need to recompute
        computedPredicates = null;
        predicatesMap = null;
        
    }


    private table_entry_bdd convert_to_bdd_entry(table_entry one_entry)
    {
        int out_label_bdd = bddengine.encodeMPLSLabel(one_entry.out_label);
        return new table_entry_bdd(out_label_bdd, one_entry.ports);
    }

    public void convert_to_bdd_table_swap()
    {
        // swap_table
        for(Integer one_in_label : swap_table.keySet())
        {

            int in_label_bdd = bddengine.encodeMPLSLabel(one_in_label);
            table_entry_bdd entry_bdd = convert_to_bdd_entry(swap_table.get(one_in_label));
            swap_table_bdd.put(in_label_bdd, entry_bdd);

        }

    }

    public void convert_to_bdd_table()
    {
        int entry_num = 0;
        // swap_table
        for(Integer one_in_label : swap_table.keySet())
        {

            int in_label_bdd = bddengine.encodeMPLSLabel(one_in_label);
            table_entry_bdd entry_bdd = convert_to_bdd_entry(swap_table.get(one_in_label));
            swap_table_bdd.put(in_label_bdd, entry_bdd);
            
            entry_num ++;

        }

        // entrance_table
        all_encapsulated = bddengine.BDDFalse;
        for(ForwardingRule fec : entrance_table.keySet())
        {
            int fec_bdd = bddengine.encodeDstIPPrefix(fec.getdestip(), fec.getprefixlen());
            table_entry_bdd entry_bdd = convert_to_bdd_entry(entrance_table.get(fec));
            entrance_table_bdd.put(fec_bdd, entry_bdd);
            
            all_encapsulated = bddengine.getBDD().orTo(all_encapsulated, fec_bdd);
            
            entry_num ++;
        }
        
        System.out.println(entry_num + " entries");
        
        
    }

    public void read_mpls_table(File infile)
    {
        Scanner OneLine = null;
        try {
            OneLine = new Scanner (infile);
            OneLine.useDelimiter("\n");
            //scanner.useDelimiter(System.getProperty("line.separator"));
            // doesn't work for .conf files
        } catch (FileNotFoundException e) {
            System.out.println (infile.getName() + " not found!"); // for debugging
            System.exit (0); // Stop program if no file found
        }

        while(OneLine.hasNext())
        {
            String linestr = OneLine.next();
            if(linestr.length() > 0)
            {
                String[] tokens = linestr.split(" ");
                int in_label = Integer.parseInt(tokens[0]);
                HashSet<String> ports = new HashSet<String>();
                for(int i = 3; i < tokens.length; i ++)
                {
                    if(BuildNetwork.is_al2s)
                    {
                        ports.add(tokens[i]);
                    }else
                    {
                        // try to remove '.'
                        ports.add(tokens[i].split("\\.")[0]);
                    }
                }
                if(tokens[1].equals("Pop"))
                {
                    //add_pop_entry(in_label, ports);
                    add_swap_entry(in_label, 3, ports);
                }else if(tokens[1].equals("Swap"))
                {
                    int out_label = Integer.parseInt(tokens[2]);
                    add_swap_entry(in_label, out_label, ports);
                }
            }
        }

        convert_to_bdd_table_swap();
    }

    /**
     * 
     * @param in_pkts
     * use transformationMap
     * @return
     */
    public HashMap<String, FWDAPSet> ForwardAction(FWDAPSet in_pkts)
    {
        HashMap<String, FWDAPSet> fwded = new HashMap<String, FWDAPSet> ();

        for(int ap : transformationMap.keySet())
        {
            if(in_pkts.contains(ap))
            {
                HashMap<String, FWDAPSet> afters = transformationMap.get(ap);
                for(String port : afters.keySet())
                {
                    if(fwded.containsKey(port))
                    {
                        fwded.get(port).union(afters.get(port));
                    }else
                    {
                        fwded.put(port, new FWDAPSet(afters.get(port)));
                    }
                }
            }
        }

        return fwded;
    }

    /**
     * assume that all input packets have mpls labels
     * for label switching
     * @param port
     * @param in_pkt
     * @return
     */

    public HashMap <String, Integer> ForwardAction(int in_pkt)
    {
        HashMap <String, Integer> fwded = new HashMap<String, Integer>();

        /*
		int mpls_pkt = bddengine.separateMPLS(in_pkt);
		if(mpls_pkt == BDDACLWrapper.BDDFalse)
		{
			return fwded;
		}*/

        int poped_pkt = bddengine.popMPLSLabel(in_pkt, reserveLabels_bdd);
        if(poped_pkt != BDDACLWrapper.BDDFalse)
        {
            fwded.put(recv_port, poped_pkt);
        }


        BDD theBDD = bddengine.getBDD();

        for(Integer in_label_bdd : swap_table_bdd.keySet())
        {

            HashSet<String> ports = swap_table_bdd.get(in_label_bdd).ports;
            int out_label_bdd = swap_table_bdd.get(in_label_bdd).out_label_bdd;
            int out_pkt = bddengine.swapMPLSLabel(in_pkt, in_label_bdd, out_label_bdd);
            if(out_pkt == BDDACLWrapper.BDDFalse)
            {
                continue;
            }

            for(String one_port : ports)
            {
                if(fwded.containsKey(one_port))
                {
                    int old_pkt = fwded.get(one_port);
                    fwded.put(one_port, theBDD.ref(theBDD.or(old_pkt, out_pkt)));
                    theBDD.deref(old_pkt);
                }else
                {
                    fwded.put(one_port, out_pkt);
                    theBDD.ref(out_pkt);
                }
            }
            // out_pkt already ref once in swap mpls label
            theBDD.deref(out_pkt);

            //break;

        }

        return fwded;
    }

    /**
     * 
     * @param in_pkts
     * use transformationMap
     * remove packets from in_pkts that are encapsulated
     * @return
     */
    public HashMap<String, FWDAPSet> EncapsulateAction(FWDAPSet in_pkts)
    {

        HashMap<String, FWDAPSet> fwded = new HashMap<String, FWDAPSet> ();

        FWDAPSet encapsulated = new FWDAPSet(BDDACLWrapper.BDDFalse);

        for(int ap : transformationMap.keySet())
        {
            if(in_pkts.contains(ap))
            {
                HashMap<String, FWDAPSet> afters = transformationMap.get(ap);
                for(String port : afters.keySet())
                {
                    if(fwded.containsKey(port))
                    {
                        fwded.get(port).union(afters.get(port));
                    }else
                    {
                        fwded.put(port, new FWDAPSet(afters.get(port)));
                    }
                }
                //encapsulated.
                encapsulated.add(ap);
            }
        }

        // remove packets that are encapsulated
        in_pkts.diff(encapsulated);

        return fwded;

    }

    /**
     * assume the input are non mpls labeled
     * @param in_pkt
     * @return
     */
    public HashMap<String, Integer> EncapsulateAction(int in_pkt)
    {
        HashMap <String, Integer> fwded = new HashMap<String, Integer>();
        /**
         * TO DO
         */
        BDD theBDD = bddengine.getBDD();

        for(Integer fec : entrance_table_bdd.keySet())
        {

            HashSet<String> ports = entrance_table_bdd.get(fec).ports;
            int out_label_bdd = entrance_table_bdd.get(fec).out_label_bdd;
            int out_pkt_no_label = theBDD.ref(theBDD.and(in_pkt, fec));
            if(out_pkt_no_label == BDDACLWrapper.BDDFalse)
            {
                continue;
            }

            int out_pkt = bddengine.pushMPLSLabel(out_pkt_no_label, out_label_bdd);
            theBDD.deref(out_pkt_no_label);

            for(String one_port : ports)
            {
                if(fwded.containsKey(one_port))
                {
                    int old_pkt = fwded.get(one_port);
                    fwded.put(one_port, theBDD.ref(theBDD.or(old_pkt, out_pkt)));
                    theBDD.deref(old_pkt);
                }else
                {
                    fwded.put(one_port, out_pkt);
                    theBDD.ref(out_pkt);
                }
            }
            // out_pkt already ref once in swap mpls label
            theBDD.deref(out_pkt);

            //break;

        }

        return fwded;
    }

    public static void main(String[] args) throws IOException
    {
        MPLSTable mt = new MPLSTable();
        MPLSTable.set_bdd_engine(new BDDACLWrapper());

        mt.read_mpls_table(new File("i2new/atlampls"));

        System.out.println(mt.swap_table_bdd.size());

    }

}